#!python

import math
import random
import xml.dom.minidom as minidom


SKELETON_FILE = "skeleton.plots"
MIN_DISTANCE = 3
MAX_DISTANCE = 10
SHORTCUT_DISTANCE = 1
SHORTCUT_THRESHOLD = 0.9
DETOUR_DISTANCE = 15
DETOUR_THRESHOLD = 0.95


def load_skeleton():
    return minidom.parse(SKELETON_FILE)

def set_max_yield(graph, max_yield):
    keys = graph.getElementsByTagName("key")
    for key in keys:
        if key.getAttribute("attr.name") == "maxYield":
            yield_tag = graph.createElement("default")
            value = graph.createTextNode(str(max_yield))
            yield_tag.appendChild(value)
            key.appendChild(yield_tag)
            return

def gen_resource_nodes(nodes, nodes_dispersion):
    size = int(math.ceil(nodes * nodes_dispersion))
    resource_nodes = range(1, nodes + 1)
    random.shuffle(resource_nodes)
    resource_nodes = resource_nodes[0:size]
    return resource_nodes

def gen_yield(max_yield, yield_dispersion):
    value = max_yield * yield_dispersion
    value = max_yield - random.random() * value
    return int(value)

def add_nodes(graph, nodes, nodes_dispersion, yield_dispersion, max_yield):
    unbalanced_graph = nodes_dispersion < 1.0 or yield_dispersion > 0.0
    if unbalanced_graph:
        resource_nodes = gen_resource_nodes(nodes, nodes_dispersion)
    
    container = graph.getElementsByTagName("graph")[0]
    for id in range(1, nodes + 1):
        node = graph.createElement("node")
        node.setAttribute("id", "v" + str(id))
        
        if unbalanced_graph:
            node_yield = 0
            if id in resource_nodes:
                node_yield = gen_yield(max_yield, yield_dispersion)
            yield_element = graph.createElement("data")
            yield_element.setAttribute("key", "vy")
            value = graph.createTextNode(str(node_yield))
            yield_element.appendChild(value)
            node.appendChild(yield_element)
        
        container.appendChild(node)

def add_edge(graph, source, target, distance):
    container = graph.getElementsByTagName("graph")[0]
    edge = graph.createElement("edge")
    edge.setAttribute("source", "v" + str(source))
    edge.setAttribute("target", "v" + str(target))
    cost = graph.createElement("data")
    cost.setAttribute("key", "ev")
    value = graph.createTextNode(str(distance))
    cost.appendChild(value)
    edge.appendChild(cost)
    container.appendChild(edge)

def rand_distance():
    distance = int(random.random() * (MAX_DISTANCE - MIN_DISTANCE)) + MIN_DISTANCE
    if random.random() > DETOUR_THRESHOLD:
        distance = DETOUR_DISTANCE
    elif random.random() > SHORTCUT_THRESHOLD:
        distance = SHORTCUT_DISTANCE
    return distance

def gen_edges(graph, nodes, edges):
    # a set of <nodes> disconnected graphs
    connected = []
    adjacency = set()
    out_degrees = []
    for node in range(0, nodes):
        connected.append(set([node]))
        out_degrees.append((node, 0))
    
    # we must connect them!
    source = 0
    target = 0
    while len(connected) > 1:
        a = random.choice(connected)
        connected.remove(a)
        b = random.choice(connected)
        connected.remove(b)
        c = a | b
        source = a.pop()
        target = b.pop()
        adjacency.add((source, target))
        adjacency.add((target, source))
        min_node = min(source, target) + 1
        max_node = max(source, target) + 1
        add_edge(graph, min_node, max_node, rand_distance())
        connected.append(c)
        out_degrees[source] = (out_degrees[source][0], out_degrees[source][1] + 1)
        out_degrees[target] = (out_degrees[target][0], out_degrees[target][1] + 1)
    connected = None # free memory
    
    # remove nodes if they were chosen every time and are connected with all others
    if out_degrees[source][1] == nodes - 1:
        out_degrees.remove(out_degrees[source])
    if out_degrees[target][1] == nodes - 1:
        out_degrees.remove(out_degrees[target])
    # 1/(<nodes>!) chance that <target> shifted one position to the left
    if out_degrees[target - 1][1] == nodes - 1:
        out_degrees.remove(out_degrees[target - 1])
    
    edges -= nodes - 1
    while edges > 0:
        a = random.choice(out_degrees)
        out_degrees.remove(a)
        b = random.choice(out_degrees)
        out_degrees.remove(b)
        if not (a[0], b[0]) in adjacency:
            adjacency.add((a[0], b[0]))
            adjacency.add((b[0], a[0]))
            min_node = min(a[0], b[0]) + 1
            max_node = max(a[0], b[0]) + 1
            add_edge(graph, min_node, max_node, rand_distance())
            if a[1] < nodes - 2:
                a = (a[0], a[1] + 1)
                out_degrees.append(a)
            if b[1] < nodes - 2:
                b = (b[0], b[1] + 1)
                out_degrees.append(b)
            edges -= 1
        else:
            out_degrees.append(a)
            out_degrees.append(b)

def gen_graph(nodes, edges, nodes_dispersion, yield_dispersion, max_yield):
    if edges < nodes - 1:
        raise RuntimeError("Not enough edges")
    if edges > nodes * (nodes - 1) / 2:
        raise RuntimeError("Too many edges")
    if nodes_dispersion <= 0.0 or nodes_dispersion > 1.0:
        raise RuntimeError("Nodes dispersion must be in (0, 1]")
    if yield_dispersion < 0.0 or yield_dispersion > 1.0:
        raise RuntimeError("Yield dispersion must be in [0, 1]")
    graph = load_skeleton()
    set_max_yield(graph, max_yield)
    add_nodes(graph, nodes, nodes_dispersion, yield_dispersion, max_yield)
    gen_edges(graph, nodes, edges)
    return graph

def dump_graph(graph, file_name):
    with open(file_name, "w") as writer:
        graph.writexml(writer=writer, addindent="    ", newl="\n", encoding="UTF-8")


if __name__ == "__main__":
    graph = gen_graph(5, 10, 1.0, 0.0, 100)
    print graph.toprettyxml(indent="    ", encoding="UTF-8")
    dump_graph(graph, "../data/test1.plots")
    print
    
    graph = gen_graph(5, 10, 0.4, 0.0, 100)
    print graph.toprettyxml(indent="    ", encoding="UTF-8")
    dump_graph(graph, "../data/test2.plots")
    print
    
    graph = gen_graph(5, 10, 1.0, 0.5, 100)
    print graph.toprettyxml(indent="    ", encoding="UTF-8")
    dump_graph(graph, "../data/test3.plots")
    print
    
    graph = gen_graph(5, 10, 0.5, 0.71, 100)
    print graph.toprettyxml(indent="    ", encoding="UTF-8")
    dump_graph(graph, "../data/test4.plots")
    print

"""
Node tag example:

    <node id="v1"/>


Edge tag example:

    <edge source="v1" target="v2">
        <data key="ev">5.0</data>
    </edge>
"""
